import numpy as np
import matplotlib.pyplot as plt
from copy import deepcopy
from random import randint
from random import uniform
from random import shuffle

from optimizer.tsp.operators import Lifecycle
from optimizer.tsp.operators import greedy_sequence, two_opt

from utils import properties

plt.rc('font', family='Times New Roman')
plt.figure(figsize=(16, 8))

city_colors = properties.palette('default')
tour_colors = properties.palette('default')


def print_edges(edges):
    for i in range(1, len(edges)):
        print(i, edges[i])


class EAX:
    """
    Edge Assembly Crossover GA for TSP
    https://doi.org/10.1287/ijoc.1120.0506
    """
    def __init__(self, data, NP: int, NC: int):
        self.__data = data
        self.__cities, self.__n = data.cities, data.n
        self.__distances, self.__cost = data.distances, data.cost
        self.__NP, self.__NC = NP, NC

        self.depot = 1
        self.population = []
        for i in range(self.__NP):
            self.population.append(self.__initialize())
        # for i in range(self.__NP):
        #     print(self.to_route(self.population[i][0]), self.population[i][1])

        self.gBest = deepcopy(min(self.population, key=lambda _: _[1]))

        self.life = Lifecycle()
        self.life.update(gBest=self.gBest[1])

    def optimize(self, iterations=None, evaluations=None, deadline=None):
        self.life.set_stopping_condition(iterations, evaluations, deadline)
        while self.life.check_stopping_condition() is False:
            shuffle(self.population)
            for i in range(self.__NP - 1):
                pA, pB = self.population[i], self.population[i + 1]
                # print_p(pA)
                # print_p(pB)

                # partition edges into AB-cycles
                self.__AB_cycles(pA, pB)

                quit(-1)

    def __initialize(self):
        route = greedy_sequence(distance_func=self.__distances, candidates=self.__cities, depot=randint(1, self.__n))
        for i in range(10):
            route = two_opt(route=route, distance_func=self.__distances)[0]
        length = self.__cost(route=route)
        edges = self.to_edges(route=route)
        return [edges, length]

    def __AB_cycles(self, pA, pB):
        EA, EB = deepcopy(pA[0]), deepcopy(pB[0])  # edges that consist of pA, pB
        AB_cycles = []  # AB-cycles to produce

        traced = np.zeros((self.__n + 1, 2), dtype=int)  # the traced edges
        source = np.zeros((self.__n + 1, 2), dtype=int)  # the source of traced edges. 1: EA or 2: EB
        vertex = -1  # the current vertex
        E = None  # reference of EA or EB in turn to trace

        traced_path = []  # record the vertices in traced edges
        while EA.sum() > 0 or EB.sum() > 0:  # EA and EB is not empty.
            if traced.sum() == 0:  # If the current traced path is empty, get a new vertex to start.
                while True:  # random select a vertex that is linked by at least one edge as the start point
                    vertex = randint(1, self.__n)  # random select a vertex
                    if any(EA[vertex]) or any(EB[vertex]):  # the vertex is linked by at least one edge
                        break

                assert vertex != -1
                traced_path.clear()
                traced_path.append(vertex)  # add the current vertex to traced path

            if E is None:  # no current edge traced
                if any(EA[vertex]) and any(EB[vertex]):  # EA and EB are both available
                    E = EA if uniform(0, 1) < 0.5 else EB  # random select EA or EB to trace an edge
                elif any(EA[vertex]):  # EA is available
                    E = EA
                elif any(EB[vertex]):  # EB is available
                    E = EB
                else:
                    print(f'Error, cannot select an edge source. E={E}')
                    assert E is not None
                    quit(-1)

            E = EA if E is EB else EB  # Switch EA and EB in turn to trace an edge. EA <-> EB

            if all(E[vertex]):  # if both of the two edges are available
                e = 0 if uniform(0, 1) < 0.5 else 1  # random trace an edge
            else:
                e = 0 if E[vertex][0] > 0 else 1  # trace the existing edge

            traced[vertex][e] = E[vertex][e]  # trace edge E[vertex][e]
            source[vertex][e] = 1 if E is EA else 2  # mark the source (EA: 0 or EB: 1) of the edge
            track = source[vertex][e]  # track of source (EA: 0 or EB: 1)

            E[vertex][e] = 0  # remove edge E[vertex][e]
            vertex = traced[vertex][e]  # switch to the new vertex

            traced_path.append(vertex)
            print(traced_path)

            if any(traced[vertex]):  # means back to a visited vertex
                for e in [0, 1]:
                    if traced[vertex][e] > 0 and source[vertex][e] != track:
                        # print('AB-cycle found.')
                        v = vertex
                        cycle = np.zeros((self.__n + 1, 2), dtype=int)  # the traced edges
                        cycle[v][e] = traced[v][e]  # record cycle
                        t = source[v][e]  # current track source
                        w = v  # current vertex
                        v = traced[v][e]  # next vertex
                        traced[w][e] = 0  # remove traced edges
                        while v != vertex:
                            for f in [0, 1]:
                                if traced[v][f] > 0 and source[v][f] != t:
                                    cycle[v][f] = traced[v][f]  # record cycle
                                    t = source[v][f]  # current track source
                                    w = v  # current vertex
                                    v = traced[v][f]  # next vertex
                                    traced[w][f] = 0  # remove traced edges
                        AB_cycles.append(cycle)

            plt.clf()
            plt.subplot(2, 3, 1)
            self.__data.draw(plot=plt, route=self.to_route(pA[0]), lc=tour_colors[1], ls='-')
            self.__data.draw(plot=plt, route=self.to_route(pB[0]), lc=tour_colors[3], ls=':')
            plt.subplot(2, 3, 2)
            self.__data.draw(plot=plt, route=self.to_route(pA[0]), lc=tour_colors[1], ls='-')
            plt.subplot(2, 3, 3)
            self.__data.draw(plot=plt, route=self.to_route(pB[0]), lc=tour_colors[3], ls=':')
            plt.subplot(2, 3, 4)
            self.__data.draw_edges(plot=plt, edges=traced, source=source, show_text=True,
                                   lc1=tour_colors[1], lc2=tour_colors[3],
                                   ls1='-', ls2=':')
            plt.subplot(2, 3, 5)
            self.__data.draw(plot=plt)
            for cycle in AB_cycles:
                self.__data.draw_edges(plot=plt, edges=cycle, source=source,
                                       lc1=tour_colors[1], lc2=tour_colors[3],
                                       ls1='-', ls2=':')
            plt.pause(0.001)
        plt.show()

    def to_edges(self, route):
        edges = np.zeros((self.__n + 1, 2), dtype=int)
        edges[route[0]][1], edges[route[-1]][0] = route[1], route[-2]
        for i in range(1, len(route) - 1):
            edges[route[i]][0], edges[route[i]][1] = route[i - 1], route[i + 1]
        return edges

    def to_route(self, edges, depot=None):
        route = [0 for _ in range(self.__n + 1)]
        if depot is None:
            route[0] = self.depot
        for i in range(len(route) - 1):
            a, b = edges[route[i]][0], edges[route[i]][1]
            if route[i + 1] == 0:
                if a not in route:
                    route[i + 1] = a
                if b not in route:
                    route[i + 1] = b
        route[-1] = self.depot
        return route


from optimizer.tsp.TSP import TSP

tsp = TSP(r'../../data/TSPLIB/eil51.tsp.txt')
eax = EAX(data=tsp, NP=10, NC=6)
eax.optimize()
